﻿//
// Created by xiaoliu on 2024/12/14.
//

#include <Windows.h>
#include <comdef.h>
#include <wbemidl.h>

#include "CertificateTool.h"
#include "ui_CertificateTool.h"
extern "C"
{
#include <openssl/applink.c>
}
#include <QDateTime>
#include <QFileInfo>
#include <QFileDialog>
#include <QMessageBox>
#include <iostream>



// 定义密钥长度
#define KEY_LENGTH 2048
// 定义公钥指数
#define PUB_EXP 3

CCertificateTool::CCertificateTool(QWidget *parent)
    : QWidget(parent)
    , m_ui(new Ui::CertificateTool)
{
    m_ui->setupUi(this);
    OpenSSL_add_all_algorithms();
}

CCertificateTool::~CCertificateTool()
{
    delete m_ui;
}

void CCertificateTool::on_btnMakeCACert_clicked()
{
    // 生成CA密钥对
    RSA* pCARsa = generateRsaKeypair();
    if (pCARsa)
    {
        // 存文件
        QString strCAPublicKeyPath = QApplication::applicationDirPath() + QString("/ca/CAPublic.pem");
        QString strCAPrivateKeyPath = QApplication::applicationDirPath() + QString("/ca/CAPrivate.pem");
        writeKeyToFile(strCAPublicKeyPath.toStdString().c_str(), strCAPrivateKeyPath.toStdString().c_str(), pCARsa);
    }
    // 创建一个X509证书对象
    X509* pX509 = X509_new();
    // 设置证书的版本号 0对应V1
    X509_set_version(pX509, 2);
    // 设置证书的序列号
    ASN1_INTEGER_set(X509_get_serialNumber(pX509), 1);
    // 设置证书的有效期
    X509_gmtime_adj(X509_get_notBefore(pX509), 0);
    X509_gmtime_adj(X509_get_notAfter(pX509), 1577836799L);
    // 设置公钥
    EVP_PKEY *pPublickey = EVP_PKEY_new();
    EVP_PKEY_assign_RSA(pPublickey, pCARsa);
    X509_set_pubkey(pX509, pPublickey);

    std::string strC = m_ui->editCA_C->text().toStdString();
    std::string strS = m_ui->editCA_S->text().toStdString();
    std::string strL = m_ui->editCA_L->text().toStdString();
    std::string strO = m_ui->editCA_O->text().toStdString();
    std::string strOU = m_ui->editCA_OU->text().toStdString();
    std::string strCN = m_ui->editCA_CN->text().toStdString();

    const unsigned char* pC = (const unsigned char*)strC.c_str();
    const unsigned char* pS = (const unsigned char*)strS.c_str();
    const unsigned char* pL = (const unsigned char*)strL.c_str();
    const unsigned char* pO = (const unsigned char*)strO.c_str();
    const unsigned char* pOU = (const unsigned char*)strOU.c_str();
    const unsigned char* pCN = (const unsigned char*)strCN.c_str();

    // 颁发者信息
    X509_NAME* pIssuerName = X509_get_issuer_name(pX509);
    X509_NAME_add_entry_by_txt(pIssuerName, "C", MBSTRING_ASC, pC, -1, -1, 0);
    X509_NAME_add_entry_by_txt(pIssuerName, "S", MBSTRING_ASC, pS, -1, -1, 0);
    X509_NAME_add_entry_by_txt(pIssuerName, "L", MBSTRING_ASC, pL, -1, -1, 0);
    X509_NAME_add_entry_by_txt(pIssuerName, "O", MBSTRING_ASC, pO, -1, -1, 0);
    X509_NAME_add_entry_by_txt(pIssuerName, "OU", MBSTRING_ASC, pOU, -1, -1, 0);
    X509_NAME_add_entry_by_txt(pIssuerName, "CN", MBSTRING_ASC, pCN, -1, -1, 0);
    X509_set_issuer_name(pX509, pIssuerName);
    // 主题者信息
    X509_NAME* pSubjectName = X509_get_subject_name(pX509);
    X509_NAME_add_entry_by_txt(pSubjectName, "C", MBSTRING_ASC, pC, -1, -1, 0);
    X509_NAME_add_entry_by_txt(pSubjectName, "S", MBSTRING_ASC, pS, -1, -1, 0);
    X509_NAME_add_entry_by_txt(pSubjectName, "L", MBSTRING_ASC, pL, -1, -1, 0);
    X509_NAME_add_entry_by_txt(pSubjectName, "O", MBSTRING_ASC, pO, -1, -1, 0);
    X509_NAME_add_entry_by_txt(pSubjectName, "OU", MBSTRING_ASC, pOU, -1, -1, 0);
    X509_NAME_add_entry_by_txt(pSubjectName, "CN", MBSTRING_ASC, pCN, -1, -1, 0);
    X509_set_subject_name(pX509, pSubjectName);
    // 添加扩展：标识这是一个CA证书
    X509_EXTENSION *pX509Ex;
    X509V3_CTX ctx;
    // 初始化上下文
    X509V3_set_ctx(&ctx, pX509, pX509, nullptr, nullptr, 0);
    pX509Ex = X509V3_EXT_conf_nid(nullptr, &ctx, NID_basic_constraints, "critical,CA:TRUE");
    X509_add_ext(pX509, pX509Ex, -1);
    X509_EXTENSION_free(pX509Ex);
    // 添加Key Usage扩展，指定证书只能用于证书签署
    pX509Ex = X509V3_EXT_conf_nid(nullptr, &ctx, NID_key_usage, "critical,keyCertSign,cRLSign");
    X509_add_ext(pX509, pX509Ex, -1);
    X509_EXTENSION_free(pX509Ex);
    // 使用自己的私钥签署证书
    X509_sign(pX509, pPublickey, EVP_sha256());
    // 保存到文件
    QString strPath = QCoreApplication::applicationDirPath() + QString("/ca/ca.crt");
    FILE* pFile = fopen(strPath.toStdString().c_str(), "wb");
    PEM_write_X509(pFile, pX509);
    fclose(pFile);

    EVP_PKEY_free(pPublickey);
    RSA_free(pCARsa);
}

void CCertificateTool::on_btnMakeUserCSR_clicked()
{
    // 创建一个X509证书对象
    X509_REQ* pX509Req = X509_REQ_new();
    // 主题者信息
    std::string strC = m_ui->editUser_C->text().toStdString();
    std::string strS = m_ui->editUser_S->text().toStdString();
    std::string strL = m_ui->editUser_L->text().toStdString();
    std::string strO = m_ui->editUser_O->text().toStdString();
    std::string strOU = m_ui->editUser_OU->text().toStdString();
    std::string strCN = m_ui->editUser_CN->text().toStdString();
    const unsigned char* pC = (const unsigned char*)strC.c_str();
    const unsigned char* pS = (const unsigned char*)strS.c_str();
    const unsigned char* pL = (const unsigned char*)strL.c_str();
    const unsigned char* pO = (const unsigned char*)strO.c_str();
    const unsigned char* pOU = (const unsigned char*)strOU.c_str();
    const unsigned char* pCN = (const unsigned char*)strCN.c_str();

    X509_NAME* pSubjectName = X509_REQ_get_subject_name(pX509Req);
    X509_NAME_add_entry_by_txt(pSubjectName, "C", MBSTRING_ASC, pC, -1, -1, 0);
    X509_NAME_add_entry_by_txt(pSubjectName, "S", MBSTRING_ASC, pS, -1, -1, 0);
    X509_NAME_add_entry_by_txt(pSubjectName, "L", MBSTRING_ASC, pL, -1, -1, 0);
    X509_NAME_add_entry_by_txt(pSubjectName, "O", MBSTRING_ASC, pO, -1, -1, 0);
    X509_NAME_add_entry_by_txt(pSubjectName, "OU", MBSTRING_ASC, pOU, -1, -1, 0);
    X509_NAME_add_entry_by_txt(pSubjectName, "CN", MBSTRING_ASC, pCN, -1, -1, 0);
    X509_REQ_set_subject_name(pX509Req, pSubjectName);

    // 自定义字段
    std::string strSerialNumber = getMotherboardSerialNumber().toStdString();
    const char* pSerialNumber = strSerialNumber.c_str();
    // 创建OID对象
    ASN1_OBJECT *obj = OBJ_txt2obj("1.3.5.7.9", 1);
    if (!obj)
    {
        fprintf(stderr, "Error creating OID\n");
        return;
    }
    // 创建ASN1_STRING值
    ASN1_STRING* str = ASN1_STRING_new();
    if (!ASN1_STRING_set(str, pSerialNumber, strlen(pSerialNumber)))
    {
        fprintf(stderr, "Error setting ASN1_STRING\n");
        return;
    }
    // 创建X509_EXTENSION
    X509_EXTENSION* pEx = X509_EXTENSION_create_by_OBJ(nullptr, obj, 0, str);
    if (!pEx)
    {
        fprintf(stderr, "Error creating X509_EXTENSION\n");
        return;
    }
    // 获取现有的扩展堆栈
    STACK_OF(X509_EXTENSION) *pStackEx = X509_REQ_get_extensions(pX509Req);
    if (pStackEx == nullptr)
    {
        pStackEx = sk_X509_EXTENSION_new_null();
    }
    // 将新扩展添加到堆栈
    sk_X509_EXTENSION_push(pStackEx, pEx);
    // 添加自定义扩展到证书请求
    X509_REQ_add_extensions(pX509Req, pStackEx);
    sk_X509_EXTENSION_free(pStackEx);
    X509_EXTENSION_free(pEx);
    // 生成用户密钥对
    RSA* pUserRsa = generateRsaKeypair();
    QString strTmp = QDateTime::currentDateTime().toString("yyyyMMddHHmmss");
    if (pUserRsa)
    {
        // 存文件
        QString strCAPublicKeyPath = QApplication::applicationDirPath() + QString("/user/User_%1_Public.pem").arg(strTmp);
        QString strCAPrivateKeyPath = QApplication::applicationDirPath() + QString("/user/User_%1_Private.pem").arg(strTmp);
        writeKeyToFile(strCAPublicKeyPath.toStdString().c_str(), strCAPrivateKeyPath.toStdString().c_str(), pUserRsa);
    }
    // 设置CSR的公钥
    EVP_PKEY* pPublickey = EVP_PKEY_new();
    EVP_PKEY_assign_RSA(pPublickey, pUserRsa);
    X509_REQ_set_pubkey(pX509Req, pPublickey);
    // 创建CSR
    X509_REQ_sign(pX509Req, pPublickey, EVP_sha256());
    // 输出CSR到文件
    QString strCAPublicKeyPath = QApplication::applicationDirPath() + QString("/user/User_%1_Csr.pem").arg(strTmp);
    BIO* pBio = BIO_new_file(strCAPublicKeyPath.toStdString().c_str(), "w+");
    if (pBio)
    {
        X509_REQ_print_ex(pBio, pX509Req, XN_FLAG_RFC2253, X509_FLAG_NO_HEADER);
        PEM_write_bio_X509_REQ(pBio, pX509Req);
        BIO_free_all(pBio);
    }
    // 清理资源
    EVP_PKEY_free(pPublickey);
    X509_REQ_free(pX509Req);
}

void CCertificateTool::on_btnSign_clicked()
{
    X509* pCACert = nullptr;
    EVP_PKEY *pPrivatekey = nullptr;
    X509_REQ* pX509Req = nullptr;
    // 从文件中读取CA证书
    QString strCACertPath = m_ui->editCACert->text();
    FILE* pFileCert = fopen(strCACertPath.toStdString().c_str(), "r");
    if (pFileCert)
    {
        pCACert = PEM_read_X509(pFileCert, nullptr, nullptr, nullptr);
        fclose(pFileCert);
    }
    // 从文件中读取CA证书对应的私钥
    QString strPrivateKeyPath = m_ui->editCAPrivate->text();
    FILE* pFileKey = fopen(strPrivateKeyPath.toStdString().c_str(), "r");
    if (pFileKey)
    {
        pPrivatekey = PEM_read_PrivateKey(pFileKey, nullptr, nullptr, nullptr);
        fclose(pFileKey);
    }
    // 从文件中读取用户的CSR文件
    QString strCsrPath = m_ui->editUserCSR->text();
    FILE* pFileCsr = fopen(strCsrPath.toStdString().c_str(), "r");
    if (pFileCsr)
    {
        pX509Req = PEM_read_X509_REQ(pFileCsr, nullptr, nullptr, nullptr);
        fclose(pFileCsr);
    }

    if (pCACert && pPrivatekey && pX509Req)
    {
        // 创建一个新的X509证书对象
        X509* pUserCert = X509_new();
        // 设置证书版本
        X509_set_version(pUserCert, 2L);
        // 设置证书序列号
        ASN1_INTEGER_set(X509_get_serialNumber(pUserCert), 1);
        // 设置证书有效期
        X509_gmtime_adj(X509_get_notBefore(pUserCert), 0);
        X509_gmtime_adj(X509_get_notAfter(pUserCert), (long)60 * m_ui->spinUserCertTime->value());
        // 设置证书的签发者名称（即CA的主体名称）
        X509_set_issuer_name(pUserCert, X509_get_subject_name(pCACert));
        // 从CSR获取公钥并设置到证书
        EVP_PKEY *pUserPublicKey = X509_REQ_get_pubkey(pX509Req);
        X509_set_pubkey(pUserCert, pUserPublicKey);
        EVP_PKEY_free(pUserPublicKey);
        // 设置证书的主题名（从CSR获取）
        X509_NAME* pUserSubject = X509_REQ_get_subject_name(pX509Req);
        X509_set_subject_name(pUserCert, pUserSubject);
        // 设置扩展(从CSR获取)
        STACK_OF(X509_EXTENSION) *pStackEx = X509_REQ_get_extensions(pX509Req);
        for (int i=0;i<sk_X509_EXTENSION_num(pStackEx);++i)
        {
            X509_EXTENSION *pEx = sk_X509_EXTENSION_value(pStackEx, i);
            if (!pEx)
            {
                continue;
            }
            X509_add_ext(pUserCert, pEx, -1);
            X509_EXTENSION_free(pEx);
        }
        sk_X509_EXTENSION_free(pStackEx);
        // 签署证书
        X509_sign(pUserCert, pPrivatekey, EVP_sha256());
        // 保存证书到文件或发送给客户端
        QFileInfo info(strCsrPath);
        QString strUserCertPath = QCoreApplication::applicationDirPath() + QString("/user/%1.crt").arg(info.completeBaseName());
        strUserCertPath = strUserCertPath.replace("_Csr", "");
        FILE* pFileUserCert = fopen(strUserCertPath.toStdString().c_str(), "wb");
        PEM_write_X509(pFileUserCert, pUserCert);
        fclose(pFileUserCert);
        // 清理
        X509_free(pUserCert);
        X509_REQ_free(pX509Req);
        EVP_PKEY_free(pPrivatekey);
        X509_free(pCACert);
    }
}

void CCertificateTool::on_btnVerify_clicked()
{
    X509* pCACert = nullptr;
    X509* pUserCert = nullptr;
    // 从文件中读取CA证书
    QString strCACertPath = m_ui->editCACert0->text();
    FILE* pFileCACert = fopen(strCACertPath.toStdString().c_str(), "r");
    if (pFileCACert)
    {
        pCACert = PEM_read_X509(pFileCACert, nullptr, nullptr, nullptr);
       fclose(pFileCACert);
    }
    // 从文件中读取用户证书
    QString strUserCertPath = m_ui->editUserCert->text();
    FILE* pFileUserCert = fopen(strUserCertPath.toStdString().c_str(), "r");
    if (pFileUserCert)
    {
        pUserCert = PEM_read_X509(pFileUserCert, nullptr, nullptr, nullptr);
        fclose(pFileUserCert);
    }
    if (pCACert && pUserCert)
    {
        // 创建并初始化一个X509_STORE对象
        X509_STORE *pX509Store = X509_STORE_new();
        if (!pX509Store)
        {
            fprintf(stderr, "Failed to create X509_STORE.\n");
            return;
        }
        // 添加CA证书到X509_STORE
        if (!X509_STORE_add_cert(pX509Store, pCACert))
        {
            fprintf(stderr, "Failed to add CA certificate to pX509Store.\n");
            X509_STORE_free(pX509Store);
            return;
        }
        // 创建X509_STORE_CTX用于验证
        X509_STORE_CTX *pCtx = X509_STORE_CTX_new();
        if (!pCtx)
        {
            fprintf(stderr, "Failed to create X509_STORE_CTX.\n");
            X509_STORE_free(pX509Store);
            return;
        }

        // 初始化X509_STORE_CTX
        if (!X509_STORE_CTX_init(pCtx, pX509Store, pUserCert, nullptr))
        {
            fprintf(stderr, "Failed to initialize X509_STORE_CTX.\n");
            X509_STORE_CTX_free(pCtx);
            X509_STORE_free(pX509Store);
            return;
        }
        // 从用户证书中读取扩展
        const STACK_OF(X509_EXTENSION)* pStackEx = X509_get0_extensions(pUserCert);
        m_ui->textBrowser->clear();
        for (int i=0;i<sk_X509_EXTENSION_num(pStackEx);++i)
        {
            X509_EXTENSION* pEx = sk_X509_EXTENSION_value(pStackEx, i);
            if (!pEx)
            {
                continue;
            }
            ASN1_OBJECT* pObj = X509_EXTENSION_get_object(pEx);
            char buf[1024]{};
            OBJ_obj2txt(buf, sizeof(buf), pObj, 1);
            ASN1_OCTET_STRING* pStr = X509_EXTENSION_get_data(pEx);
            const char* pValue = reinterpret_cast<const char*>(ASN1_STRING_get0_data(pStr));
            QString strText = QString("%1: %2").arg(buf).arg(pValue);
            m_ui->textBrowser->append(strText);
        }
        // 执行验证
        int ret = X509_STORE_CTX_verify(pCtx);
        if (ret != 1)
        {
            // 如果验证失败，打印错误信息
            int error = X509_STORE_CTX_get_error(pCtx);
            QString strText = QString("Certificate verification failed: %1").arg(X509_verify_cert_error_string(error));
            QMessageBox::warning(this, QString("Error"), strText);
        }
        else
        {
            QMessageBox::information(this, QString("Succes"), QString("Certificate verified successfully"));
        }
        // 清理资源
        X509_STORE_CTX_free(pCtx);
        X509_STORE_free(pX509Store);
        X509_free(pUserCert);
        X509_free(pCACert);
    }
}

// 生成RSA密钥对
RSA* CCertificateTool::generateRsaKeypair()
{
    // 创建新的RSA对象
    RSA* pRsa = RSA_new();
    // 创建新的大数对象
    BIGNUM* pBn = BN_new();
    // 设置公钥指数
    if (!BN_set_word(pBn, PUB_EXP))
    {
        // 如果设置失败，则处理错误
        ERR_print_errors_fp(stderr);
    }
    // 生成密钥对
    if (!RSA_generate_key_ex(pRsa, KEY_LENGTH, pBn, nullptr))
    {
        // 如果生成失败，则处理错误
        ERR_print_errors_fp(stderr);
    }
    // 释放大数对象
    BN_free(pBn);
    return pRsa;
}

// 将密钥写入文件的函数
void CCertificateTool::writeKeyToFile(const char* pPublicKeyPath, const char* pPrivateKeyPath, RSA *rsa)
{
    // 写公钥到文件
    BIO *pPublicBio = BIO_new_file(pPublicKeyPath, "w+");
    PEM_write_bio_RSAPublicKey(pPublicBio, rsa);
    BIO_free_all(pPublicBio);

    // 写私钥到文件
    BIO *pPrivateBio = BIO_new_file(pPrivateKeyPath, "w+");
    PEM_write_bio_RSAPrivateKey(pPrivateBio, rsa, nullptr, nullptr, 0, nullptr, nullptr);
    BIO_free_all(pPrivateBio);
}

QString CCertificateTool::selectFile()
{
    QString strFilePath = QFileDialog::getOpenFileName(this, tr("Open File"),
                                                    "d:",
                                                    tr("Cert (*.pem *.crt *.cer)"));
    return strFilePath;
}

void CCertificateTool::on_btnCACert_clicked()
{
    QString strFilePath = selectFile();
    m_ui->editCACert->setText(strFilePath);
}

void CCertificateTool::on_btnCAPrivate_clicked()
{
    QString strFilePath = selectFile();
    m_ui->editCAPrivate->setText(strFilePath);
}

void CCertificateTool::on_btnUserCSR_clicked()
{
    QString strFilePath = selectFile();
    m_ui->editUserCSR->setText(strFilePath);
}

void CCertificateTool::on_btnCACert0_clicked()
{
    QString strFilePath = selectFile();
    m_ui->editCACert0->setText(strFilePath);
}

void CCertificateTool::on_btnUserCert_clicked()
{
    QString strFilePath = selectFile();
    m_ui->editUserCert->setText(strFilePath);
}

QString CCertificateTool::getMotherboardSerialNumber()
{
    HRESULT hres = CoInitializeEx(0, COINIT_APARTMENTTHREADED);
    if (FAILED(hres))
    {
        std::cerr << "Failed to initialize COM library. Error code = 0x" << std::hex << hres << std::endl;
        return QString();
    }

    IWbemLocator *pLoc = nullptr;
    hres = CoCreateInstance(
        CLSID_WbemLocator,
        0,
        CLSCTX_INPROC_SERVER,
        IID_IWbemLocator, (LPVOID *)&pLoc);

    if (FAILED(hres))
    {
        std::cerr << "Failed to create IWbemLocator object. "
                  << "Error code = 0x"
                  << std::hex << hres << std::endl;
        CoUninitialize();
        return QString();
    }

    IWbemServices *pSvc = nullptr;
    hres = pLoc->ConnectServer(
        _bstr_t(L"ROOT\\CIMV2"), // WMI namespace
        NULL,                    // User name
        NULL,                    // User password
        0,                       // Locale
        NULL,                    // Security flags
        0,                       // Authority
        0,                       // Context object
        &pSvc                    // IWbemServices proxy
    );

    if (FAILED(hres))
    {
        std::cerr << "Could not connect. Error code = 0x" << std::hex << hres << std::endl;
        pLoc->Release();
        CoUninitialize();
        return QString();
    }

    hres = CoSetProxyBlanket(
        pSvc,                        // Indicates the proxy to set
        RPC_C_AUTHN_WINNT,           // RPC_C_AUTHN_xxx
        RPC_C_AUTHZ_NONE,            // RPC_C_AUTHZ_xxx
        nullptr,                        // Server principal name
        RPC_C_AUTHN_LEVEL_CALL,      // RPC_C_AUTHN_LEVEL_xxx
        RPC_C_IMP_LEVEL_IMPERSONATE, // RPC_C_IMP_LEVEL_xxx
        nullptr,                        // Client identity
        EOAC_NONE                    // Proxy capabilities
    );

    if (FAILED(hres))
    {
        std::cerr << "Could not set proxy blanket. Error code = 0x"
                  << std::hex << hres << std::endl;
        pSvc->Release();
        pLoc->Release();
        CoUninitialize();
        return QString();
    }

    IEnumWbemClassObject* pEnumerator = nullptr;
    hres = pSvc->ExecQuery(
        bstr_t("WQL"),
        bstr_t("SELECT SerialNumber FROM Win32_BaseBoard"),
        WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
        NULL,
        &pEnumerator);

    if (FAILED(hres))
    {
        std::cerr << "Query for operating system caption failed."
                  << " Error code = 0x"
                  << std::hex << hres << std::endl;
        pSvc->Release();
        pLoc->Release();
        CoUninitialize();
        return QString();
    }

    // Step 7: Get the data from the query.
    IWbemClassObject *pclsObj = nullptr;
    ULONG uReturn = 0;
    QString strRes;
    while (pEnumerator)
    {
        HRESULT hr = pEnumerator->Next(WBEM_INFINITE, 1,&pclsObj, &uReturn);
        if(0 == uReturn)
        {
            break;
        }

        VARIANT vtProp;
        hr = pclsObj->Get(L"SerialNumber", 0, &vtProp, 0, 0);
        strRes = QString::fromWCharArray(vtProp.bstrVal);
        VariantClear(&vtProp);
        pclsObj->Release();
    }

    pSvc->Release();
    pLoc->Release();
    pEnumerator->Release();
    CoUninitialize();

    return strRes;
}





